import secrets
import os
from math_utils import generate_elgamal_parameters, power, mod_inverse

class ElGamal:
    """
    ElGamal 加密算法封装类。
    提供密钥生成、加密、解密，以及字符串/整数转换工具。
    """
    def __init__(self, key_size_bits=2048):
        # 生成公私钥
        self.public_key, self.private_key = self.generate_keys(key_size_bits)

    @staticmethod
    def generate_keys(key_size_bits=2048):
        """
        生成 ElGamal 密钥对。
        """
        p, q, g = generate_elgamal_parameters(key_size_bits)
        # 修改: 使用 secrets 模块安全地生成私钥 x
        # secrets.randbelow(p-2) 生成 [0, p-3] 的数, +1 后范围为 [1, p-2]
        x = secrets.randbelow(p - 2) + 1
        y = power(g, x, p)
        public_key = (p, q, g, y)
        private_key = x
        print("\nElGamal 密钥生成完毕:")
        print(f"  y = g^x mod p: {y}")
        print(f"  x (私钥): {x}")
        return public_key, private_key

    @staticmethod
    def encrypt(public_key, message_int):
        """
        使用 ElGamal 公钥加密整数消息。
        """
        p, q, g, y = public_key
        if not (0 <= message_int < p):
            raise ValueError("消息必须在 [0, p-1) 范围内")
        # 修改: 使用 secrets 模块安全地生成临时密钥 k
        k = secrets.randbelow(p - 2) + 1
        c1 = power(g, k, p)
        s = power(y, k, p)
        c2 = (message_int * s) % p
        return c1, c2

    @staticmethod
    def decrypt(private_key_x, public_params, ciphertext):
        """
        使用私钥解密 ElGamal 密文。
        """
        p, q, g = public_params
        c1, c2 = ciphertext
        s = power(c1, private_key_x, p)
        s_inv = mod_inverse(s, p)
        return (c2 * s_inv) % p

    # --- 辅助方法 ---
    @staticmethod
    def string_to_int(s):
        return int.from_bytes(s.encode('utf-8'), 'big')

    @staticmethod
    def int_to_string(i):
        length = (i.bit_length() + 7) // 8
        return i.to_bytes(length, 'big').decode('utf-8')
    
    @staticmethod
    def encrypt_bytes(public_key, data_bytes: bytes):
        p, _, _, _ = public_key
        max_len = (p.bit_length() - 1) // 8
        if len(data_bytes) >= max_len:
            raise ValueError(f"数据过长，最大允许 {max_len - 1} 字节")
        
        # 使用 os.urandom 生成的随机字节是安全的
        padded_data = data_bytes + os.urandom(1) 
        m_int = int.from_bytes(padded_data, 'big')
        return ElGamal.encrypt(public_key, m_int)

    @staticmethod
    def decrypt_bytes(private_key_x, public_params, ciphertext) -> bytes:
        padded_int = ElGamal.decrypt(private_key_x, public_params, ciphertext)
        padded_bytes = padded_int.to_bytes((padded_int.bit_length() + 7) // 8, 'big')
        return padded_bytes[:-1]